import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.colors import LogNorm
import numpy as np

# 设置中文字体
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

# 读取数据
output_dir = "results\\黑龙江2021-2023Galanaki"
lightning_with_clusters = pd.read_csv(f"{output_dir}/lightning_with_clusters.csv")
thunderstorm_metrics = pd.read_csv(f"{output_dir}/thunderstorm_metrics.csv")

# 将日期时间列转换为正确的格式
lightning_with_clusters['DATETIME'] = pd.to_datetime(lightning_with_clusters['DATETIME'])
thunderstorm_metrics['INITIATION_TIME'] = pd.to_datetime(thunderstorm_metrics['INITIATION_TIME'])
thunderstorm_metrics['CESSATION_TIME'] = pd.to_datetime(thunderstorm_metrics['CESSATION_TIME'])

# 计算每个雷暴云团的质心
storm_centroids = lightning_with_clusters.groupby('THUNDERSTORM_ID').agg({
    'LONGITUDE': 'mean',
    'LATITUDE': 'mean'
}).reset_index()
storm_centroids.columns = ['THUNDERSTORM_ID', 'CENTROID_LON', 'CENTROID_LAT']

# 合并质心数据到雷暴云团指标中
thunderstorm_metrics = pd.merge(
    thunderstorm_metrics,
    storm_centroids,
    on='THUNDERSTORM_ID',
    how='left'
)

# 创建一个包含所有可视化的画布
fig = plt.figure(figsize=(18, 15))

# 可视化 1：闪电聚类的空间分布（使用采样以减少绘图负荷）
ax1 = fig.add_subplot(2, 2, 1)
sample_size = min(10000, len(lightning_with_clusters))  # 限制样本量以避免绘图过慢
sample_indices = np.random.choice(len(lightning_with_clusters), sample_size, replace=False)
sns.scatterplot(
    data=lightning_with_clusters.iloc[sample_indices],
    x='LONGITUDE',
    y='LATITUDE',
    hue='LIGHTNING_CLUSTER',
    palette='viridis',
    s=10,
    ax=ax1
)
ax1.set_title('闪电聚类空间分布')
ax1.set_xlabel('经度')
ax1.set_ylabel('纬度')
ax1.legend(title='聚类ID', loc='upper right', bbox_to_anchor=(1.3, 1))

# 可视化 2：雷暴云团的空间分布（按强度着色）
ax2 = fig.add_subplot(2, 2, 2)
sns.scatterplot(
    data=thunderstorm_metrics,
    x='CENTROID_LON',
    y='CENTROID_LAT',
    hue='TOTAL_FLASHES',
    size='AREA_KM2',
    sizes=(50, 300),
    palette='plasma',
    alpha=0.7,
    ax=ax2
)
ax2.set_title('雷暴云团空间分布（大小=面积，颜色=总闪电次数）')
ax2.set_xlabel('经度')
ax2.set_ylabel('纬度')
ax2.legend(title='雷暴云团', loc='upper right', bbox_to_anchor=(1.3, 1))

# 可视化 3：雷暴云团持续时间与闪电频率关系
ax3 = fig.add_subplot(2, 2, 3)
sns.scatterplot(
    data=thunderstorm_metrics,
    x='DURATION_MINUTES',
    y='FLASHES_PER_MINUTE',
    hue='SPEED_KM_H',
    palette='coolwarm',
    ax=ax3
)
ax3.set_title('雷暴云团持续时间与闪电频率关系')
ax3.set_xlabel('持续时间（分钟）')
ax3.set_ylabel('每分钟闪电次数')
ax3.legend(title='移动速度 (km/h)', loc='upper right')

# 可视化 4：雷暴云团强度指标分布
ax4 = fig.add_subplot(2, 2, 4)
metrics_to_plot = ['TOTAL_FLASHES', 'DURATION_MINUTES', 'SPEED_KM_H', 'AREA_KM2']
thunderstorm_metrics[metrics_to_plot].plot(kind='box', ax=ax4)
ax4.set_title('雷暴云团强度指标分布')
ax4.set_ylabel('值')
plt.tight_layout()
plt.show()

# 可视化 5：雷暴云团时间分布（随季节）
plt.figure(figsize=(12, 6))
lightning_with_clusters['MONTH'] = lightning_with_clusters['DATETIME'].dt.month
monthly_flashes = lightning_with_clusters.groupby('MONTH')['LIGHTNING_CLUSTER'].count().reset_index()
sns.barplot(data=monthly_flashes, x='MONTH', y='LIGHTNING_CLUSTER')
plt.title('每月闪电次数分布')
plt.xlabel('月份')
plt.ylabel('闪电次数')
plt.show()

# 可视化 6：雷暴云团移动方向分布
plt.figure(figsize=(10, 10))
# 过滤掉无效的方向值
valid_bearings = thunderstorm_metrics[thunderstorm_metrics['BEARING_DEGREES'] >= 0]['BEARING_DEGREES']
n_bins = 16
hist, bins = np.histogram(valid_bearings, bins=n_bins, range=(0, 360))
width = 2 * np.pi / n_bins
ax = plt.subplot(111, projection='polar')
theta = np.linspace(0.0, 2 * np.pi, n_bins, endpoint=False)
bars = ax.bar(theta, hist, width=width, bottom=0.0)
ax.set_theta_zero_location("N")  # 北方为0度
ax.set_theta_direction(-1)  # 顺时针方向为正
ax.set_title("雷暴云团移动方向分布")
plt.show()